package com.zjh.common.datasource;

import com.alibaba.druid.pool.DruidDataSource;
import com.alibaba.druid.spring.boot.autoconfigure.properties.DruidStatProperties;
import com.alibaba.druid.spring.boot.autoconfigure.stat.DruidFilterConfiguration;
import com.alibaba.druid.spring.boot.autoconfigure.stat.DruidSpringAopConfiguration;
import com.alibaba.druid.spring.boot.autoconfigure.stat.DruidStatViewServletConfiguration;
import com.alibaba.druid.spring.boot.autoconfigure.stat.DruidWebStatFilterConfiguration;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.AutoConfigureBefore;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.autoconfigure.jdbc.DataSourceProperties;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.ApplicationContext;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Import;
import org.springframework.context.annotation.Primary;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;

/**
 * druid 数据源
 *
 * @author shanhuiming
 */
@Configuration(proxyBeanMethods = false)
@AutoConfigureBefore(DataSourceAutoConfiguration.class)
@EnableConfigurationProperties({DruidStatProperties.class, DynamicDataSourceProperties.class})
@Import({DruidSpringAopConfiguration.class, DruidStatViewServletConfiguration.class,
        DruidWebStatFilterConfiguration.class, DruidFilterConfiguration.class})
public class DynamicDataSourceConfiguration {

    @Autowired
    private DataSourceProperties dataSourceProperties;

    @Autowired
    private DynamicDataSourceProperties dynamicDataSourceProperties;

    @Autowired
    private ApplicationContext applicationContext;

    @Bean
    @Primary
    public DataSource dataSource() throws Exception {
        Map<String, DataSourceProperties> propertiesMap = dynamicDataSourceProperties.getDynamic();
        if (propertiesMap != null && !propertiesMap.isEmpty()) {
            DruidDataSource primary = null;
            Map<Object, Object> dataSourceMap = new HashMap<>();
            for (Map.Entry<String, DataSourceProperties> entry : propertiesMap.entrySet()) {
                DruidDataSourceWrapper dataSource = applicationContext.getBean(DruidDataSourceWrapper.class, entry.getValue());
                dataSource.afterPropertiesSet();
                dataSource.init();
                dataSourceMap.put(entry.getKey(), dataSource);
                if (Objects.equals(entry.getKey(), "primary")) {
                    primary = dataSource;
                }
            }
            return new DynamicDataSource(primary, dataSourceMap);
        } else {
            DruidDataSourceWrapper dataSource = applicationContext.getBean(DruidDataSourceWrapper.class, dataSourceProperties);
            dataSource.afterPropertiesSet();
            dataSource.init();

            // 虽然只有一个，还是跟上面保持统一，方便后面处理
            Map<Object, Object> dataSourceMap = new HashMap<>();
            dataSourceMap.put("primary", dataSource);
            return new DynamicDataSource(dataSource, dataSourceMap);
        }
    }
}
